import torch
import cv2
import os
import xml.etree.ElementTree as ET


def corners_to_center(x1, y1, x2, y2):
    """
    将左上角和右下角坐标转换成中心坐标和宽高
    """
    center_x = (x1 + x2) / 2
    center_y = (y1 + y2) / 2
    width = x2 - x1
    height = y2 - y1
    return [center_x, center_y, width, height]


def voc2yolo(root_path, classes, style="train"):
    features = []
    ground_truths = []
    cls_list = []
    for i, class_name in enumerate(classes):
        cls = torch.zeros((len(classes) + 1,))
        # 生成训练数据集
        cls_train_path = os.path.join(root_path, 'ImageSets/Main', f"{class_name}_{style}.txt")
        with open(cls_train_path, 'r') as f:
            train_list = f.readlines()
        for train_item in train_list:
            fname, label = train_item.strip().split()
            label = int(label)
            image = cv2.imread(os.path.join(root_path, 'JPEGImages', fname + '.jpg'))
            h, w = image.shape[:2]
            xml_data = ET.parse(os.path.join(root_path, 'Annotations', fname + '.xml'))
            image = cv2.resize(image, (448, 448))
            image = torch.from_numpy(image).permute([2, 0, 1])
            if label < 0:
                continue  # 除去不包含该类别的图片
            cls[i + 1] = 1
            for obj in xml_data.findall('object'):
                cls_name = obj.find('name').text
                if cls_name != class_name:  # 不是该类别
                    continue
                bndbox = obj.find('bndbox')
                xmin = float(bndbox.find('xmin').text) / w
                ymin = float(bndbox.find('ymin').text) / h
                xmax = float(bndbox.find('xmax').text) / w
                ymax = float(bndbox.find('ymax').text) / h
                center = corners_to_center(xmin, ymin, xmax, ymax)
                features.append(image)
                ground_truths.append(torch.tensor(center))
                cls_list.append(cls)

    return {"features": features, "ground_truths": ground_truths, "cls_list": cls_list}


if __name__ == '__main__':
    classes = ['cat', 'dog']
    root_path = "../VOCdevkit/VOC2007"
    train_data = voc2yolo(root_path, classes, style="train")
    torch.save(train_data, "../VOCdevkit/cat_dog_train_data.pth")
    valid_data = voc2yolo(root_path, classes, style="val")
    torch.save(valid_data, "../VOCdevkit/cat_dog_valid_data.pth")
